# -*- encoding: utf-8 -*-

"""
@File:      tc_lmdb_lmdb_func_004.py
@Time:      2024/03/28 14:33:20
@Author:    chenchunhu
@Version:   1.0
@Contact:   wb-cch358909@alibaba-inc.com
@License:   Mulan PSL v2
@Modify:    chenchunhu
"""

from common.basetest import LocalTest

class Test(LocalTest):
    """
    See tc_lmdb_lmdb_func_004.yaml for details

    :avocado: tags=P2,noarch,local,fix
    """

    PARAM_DIC = {"pkg_name": "lmdb lmdb-devel gcc"}
    def setUp(self):
        super().setUp(self.PARAM_DIC)
        cmdline = '''cat >lmdb_test4.c<<"EOF"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <lmdb.h>

#define ENV_PATH "./lmdb_test4_db"
#define MAX_SIZE 1024
#define NUM_ITEMS 10

// 定义宏进行测试控制
#define STEP_WRITE 1
#define STEP_CRASH 2
#define STEP_READ 3


void check(int rc) {
    if (rc != MDB_SUCCESS) {
        fprintf(stderr, "LMDB Error: %s\\n", mdb_strerror(rc));
        exit(EXIT_FAILURE);
    }
}

int main() {
    MDB_env *env;
    MDB_dbi dbi;
    MDB_txn *txn;
    MDB_val key, data;
    char value[MAX_SIZE];
    int rc;
    int i;

    // 设置随机数种子
    srand((unsigned)time(NULL));

    // 创建或打开一个LMDB环境
    rc = mdb_env_create(&env);
    check(rc);
    rc = mdb_env_set_mapsize(env, MAX_SIZE * NUM_ITEMS * 4);
    check(rc);
    rc = mdb_env_open(env, ENV_PATH, MDB_CREATE, 0664);
    check(rc);

    switch (TEST_STEP) {
        case STEP_WRITE:
            // 写入阶段
            rc = mdb_txn_begin(env, NULL, 0, &txn);  // 开始事务
            check(rc);
            rc = mdb_dbi_open(txn, NULL, MDB_CREATE, &dbi);  // 打开数据库
            check(rc);

            for (i = 0; i < NUM_ITEMS; i++) {
                key.mv_size = sizeof(int);
                key.mv_data = &i;
                snprintf(value, sizeof(value), "Item #%d", i);
                data.mv_size = strlen(value) + 1;
                data.mv_data = value;
                rc = mdb_put(txn, dbi, &key, &data, 0);  // 写入数据
                check(rc);
            }

            rc = mdb_txn_commit(txn);  // 提交事务
            check(rc);
            
            mdb_dbi_close(env, dbi);  // 关闭数据库句柄
            printf("Data write finished.\\n");
            break;

        case STEP_CRASH:
            // 模拟崩溃的写入
            rc = mdb_txn_begin(env, NULL, 0, &txn);
            check(rc);
            rc = mdb_dbi_open(txn, NULL, MDB_CREATE, &dbi);
            check(rc);
            for (i = 0; i < NUM_ITEMS; i++) {
                if (i == 5) {
                    // 模拟崩溃发生在第5个键写入之前
                    printf("Simulating crash before writing Key %d...\\n", i);
                    mdb_txn_abort(txn);  // 放弃事务模拟崩溃
                    mdb_dbi_close(env, dbi);
                    mdb_env_close(env);
                    exit(EXIT_FAILURE); // 实际上我们不应该到这儿，退出来模拟崩溃
                }
                key.mv_size = sizeof(int);
                key.mv_data = &i;
                snprintf(value, MAX_SIZE, "Item #%d", i);
                data.mv_size = strlen(value) + 1;
                data.mv_data = value;
                rc = mdb_put(txn, dbi, &key, &data, 0);
                check(rc);
            }
            // 如果没有模拟崩溃，提交事务
            rc = mdb_txn_commit(txn);
            check(rc);
            mdb_dbi_close(env, dbi);
            break;

        case STEP_READ:
            // 数据恢复与验证
            rc = mdb_txn_begin(env, NULL, MDB_RDONLY, &txn);
            check(rc);
            rc = mdb_dbi_open(txn, NULL, 0, &dbi);
            check(rc);

            for (i = 0; i < NUM_ITEMS; i++) {
                key.mv_size = sizeof(int);
                key.mv_data = &i;
                rc = mdb_get(txn, dbi, &key, &data);
                if (rc == MDB_NOTFOUND) {
                    printf("Key %d not found (possibly due to simulated crash)\\n", i);
                } else {
                    check(rc);  // 确保没有其他错误
                    printf("Recovered Key %d: %s\\n", i, (char *)data.mv_data);
                }
            }

            mdb_txn_abort(txn);  // 结束读取事务
            mdb_dbi_close(env, dbi);
            printf("Data recovery finished.\\n");
            break;

        default:
            printf("Invalid test step.\\n");
    }

    mdb_env_close(env);
    return EXIT_SUCCESS;
}
EOF'''
        self.cmd(cmdline)
        self.cmd("mkdir -p lmdb_test4_db")


    def test(self):
        self.cmd("gcc -o lmdb_test4 lmdb_test4.c -llmdb -DTEST_STEP=STEP_WRITE")
        code, lmdb_result = self.cmd("./lmdb_test4")
        self.assertIn("Data write finished.", lmdb_result)
        self.cmd("gcc -o lmdb_test4 lmdb_test4.c -llmdb -DTEST_STEP=STEP_CRASH")
        self.cmd("./lmdb_test4 > lmdb.log", ignore_status=True)
        code, lmdb_result = self.cmd("cat lmdb.log", ignore_status=True)
        self.assertIn("Simulating crash before writing Key 5", lmdb_result)
        self.cmd("gcc -o lmdb_test4 lmdb_test4.c -llmdb -DTEST_STEP=STEP_READ")
        code, lmdb_result = self.cmd("./lmdb_test4")
        self.assertIn("Recovered Key 0: Item #0", lmdb_result)

    def tearDown(self):
        super().tearDown(self.PARAM_DIC)
        self.cmd("rm -rf lmdb_test4.c lmdb_test4 lmdb.log lmdb_test4_db")
